# -*- coding: utf-8 -*-
"""
Created on Mon Sep 12 21:33:00 2022

@author: Runx_cao
"""
import numpy as np 
import matplotlib.pyplot as plt
import pandas as pd
import Stochastic_Rainfall_v2 as sr
import math
import settings


sr.resolution = 100 # size of the grid.
sr.ws = 40000 #storm envelope width.
sr.ls = 40000 #storm envelope length
storm_envelope = np.zeros((sr.ls//sr.resolution, sr.ws//sr.resolution)) #size of storm envelope

grid_n = sr.ws//sr.resolution
grid_m = sr.ls//sr.resolution
'''
把storm cells 的信息都写到一个DataFrame里面
 x  y   r_max   s1  s2
 
'''
sr.num = sr.num_gen() # num of storm cells,需要改sr中的num，初始设置是10000
num = sr.num
pos = sr.loc_gen(num) # storm cells 的中心点坐标, row number and col number
rmax = sr.r_gen()
s1 = sr.s1_gen()
s2 = sr.s1_gen()
pos["r_max"] = rmax
pos["s1"] = s1
pos["s2"] = s2

def neighbor_search(x, y, s1, s2):
    '''
    Storm cells is ellipse，长短半轴分别是S1与S2.
    搜索策略是在每一个长以2×s1与2×s2为边长的长方形内部，
    '''
    neighbor = []
    dist1 = round((s1*1000/sr.resolution-1))
    dist2 = round((s2*1000/sr.resolution-1))
    x_start = max(x-dist1, 0)
    y_start = max(y-dist2, 0)
    x_end = min(x + dist1, grid_m -1)
    y_end = min(y + dist2, grid_n-1)
    for i in range(x_start, x_end):
        for j in range(y_start, y_end):
            ellipse = ((i-x)*sr.resolution/1000)**2/s1**2 + \
                      ((j-y)*sr.resolution/1000)**2/s2**2
            if ellipse>1:
                pass
            elif i!=x or j!=y:
                neighbor.append((i,j))
    return(neighbor)

def storm_disc():
    '''
    draw the strom disc in strom_envelope
    '''
    for i in range(num):
        # storm cell's intensity
        r_int = pos.iloc[i].r_max
        # setting the rainfall instensity in the storm cell cneter
        storm_envelope[int(pos.iloc[i].x), int(pos.iloc[i].y)] += r_int
        # the idx of neighbor cells
        nb_idx = neighbor_search(int(pos.iloc[i].x), int(pos.iloc[i].y), pos.iloc[i].s1, pos.iloc[i].s2)
        '''
        for idx in nb_idx:
            storm_envelope[idx] += r_int 
        '''
    #return
    return(nb_idx)

def storm_spread(idx, x0, y0, s1, s2, r_int):
    '''
    rainfall cells are described by a Gaussian function
    idx : index of the grid in the storm cell, (x, y) 
    '''
    #dist_x = abs(idx[0]-x0)*sr.resolution
    #dist_y = abs(idx[1]-y0)*sr.resolution
    #r = r_int/math.sqrt(2*math.pi*s1**2)*\
    #        math.exp(-dist_x**2*dist_y**2/(2*s1**2*s2**2))
            
    dist1 = (idx[0]-x0)**2
    dist2 = (idx[1]-y0)**2
    tmp = -dist1/(2*(round((s1*1000/sr.resolution-1)))**2)\
          -dist2/(2*(round((s1*1000/sr.resolution-1)))**2)
    r = r_int*math.exp(tmp)
    return(r)
        
                
# storm_disc()
'''
storm_envelope[int(pos.iloc[0].x), int(pos.iloc[0].y)] += r_int
r_int = pos.iloc[0].r_max
nb_idx = neighbor_search(int(pos.iloc[0].x), int(pos.iloc[0].y), pos.iloc[0].s1, pos.iloc[0].s2)
for idx in nb_idx:
    storm_envelope[idx] += r_int 
'''

r_int = pos.iloc[0].r_max
# setting the rainfall instensity in the storm cell cneter
storm_envelope[int(pos.iloc[0].x), int(pos.iloc[0].y)] += r_int
# the idx of neighbor cells
nb_idx = neighbor_search(int(pos.iloc[0].x), int(pos.iloc[0].y), pos.iloc[0].s1, pos.iloc[0].s2)
for idx in nb_idx:
    storm_envelope[idx] += storm_spread(idx, pos.iloc[0].x, pos.iloc[0].y, pos.iloc[0].s1, pos.iloc[0].s2, pos.iloc[0].r_max)

plt.imshow(storm_envelope)